package com.xp.ai.ragdemo;

import com.xp.ai.util.ApiKey;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;

import java.io.File;
import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

/**
 *
 * rag demo
 * rag 前置步骤，知识库文档的解析与录入
 * @author xp
 */
public class RagDocmentParse {

    final static String FILE_PATH ="D:\\git_xp\\data-meta-by-ai\\LocalOllamaLanchat4j\\src\\main\\resources\\meituan-question.txt";
    public static void main(String[] args) throws URISyntaxException {
        File file = new File(FILE_PATH);
        Path path = Paths.get(file.toURI());
        //文件解析器
        TextDocumentParser textDocumentParser = new TextDocumentParser();
        //指定文件解析器和 文件加载器 获取文件
        Document document = FileSystemDocumentLoader.loadDocument(path,textDocumentParser);
        //还需要指定一个拆分器
        //新建一个拆分器
        MySplitter mySplitter = new MySplitter();
        DocumentSplitter splitter = DocumentSplitters.recursive(150, 10);
//        DocumentByLineSplitter documentByLineSplitter = new DocumentByLineSplitter(150, 10);
        //使用拆分器拆分文档
//        List<TextSegment> textSegments = documentByLineSplitter.split(document);
        //将文档拆分为一段一段的
        List<TextSegment> textSegments1 = splitter.split(document);
//        List<TextSegment> textSegments = mySplitter.split(document);
//        System.out.println(textSegments);
        System.out.println(textSegments1);
        //将所有的文本向量化，并且存在向量数据库中
        //向量数据库
        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                .host("localhost")
                .port(5433)
                .user("root")
                .password("123456")
                .database("postgres")
                .table("question_embedding")
                .dimension(1024)
                .useIndex(true)
                .indexListSize(100)
                .createTable(true)
                .dropTableFirst(false)
                .build();

        //创建向量化的模型
        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ApiKey.GJ_BASE_URL)
                .apiKey(ApiKey.GJ_API_KEY)
                .modelName(ApiKey.GJ_EMBEDDING_MODEL)
                .build();
        //将所有的文件向量化
        Response<List<Embedding>> embeddingList = embeddingModel.embedAll(textSegments1);
        //将向量存储到向量数据库中
        List<Embedding> content = embeddingList.content();
        //存储到向量数据库
        embeddingStore.addAll(content, textSegments1);

    }
}
